import torch
from torch import nn
from torchsummary import summary

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.c1=nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)#默认stride是1，填充是0
        self.sig=nn.Sigmoid()#sig只是定义，可多次调用
        self.s2=nn.AvgPool2d(kernel_size=2, stride=2)#AlexNet之前只有平均池化，这是定义池化层
        self.c3=nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, padding=0)#padding默认是0
        self.s4=nn.AvgPool2d(kernel_size=2, stride=2)

        self.flatten=nn.Flatten()#定义平展层
        self.f5=nn.Linear(in_features=5*5*16, out_features=120)#定义线性全连接层 
        self.f6=nn.Linear(in_features=120, out_features=84)
        self.f7=nn.Linear(in_features=84, out_features=10)

    def forward(self, x):
        x=self.sig(self.c1(x))
        x=self.s2(x)
        x=self.sig(self.c3(x))
        x=self.s4(x)
        x=self.flatten(x)#平展
        x=self.f5(x)
        x=self.f6(x)
        x=self.f7(x)
        return x

"""
1.直接运行：
当你直接运行一个 Python 脚本时（例如通过命令行 python script.py），Python 解释器会将该脚本的 __name__ 变量设置为 '__main__'。
因此，if __name__ == '__main__': 这一行代码会执行，通常用于执行脚本的主要功能或测试代码。
2.作为模块导入：
当你将一个 Python 文件作为模块导入到另一个文件中时（例如 import script），Python 解释器会将该文件的 __name__ 变量设置为文件名（不包括 .py 扩展名）。
在这种情况下，if __name__ == '__main__': 这一行代码不会执行，因为 __name__ 不等于 '__main__'。
"""
if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(device)#cpu

    model = LeNet().to(device)#实例化，模型放到设备里面
    print(summary(model, (1,28,28)))

